from PyQt5.QtWidgets import QOpenGLWidget, QSizePolicy
from PyQt5.QtCore import Qt
from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import numpy as np

class PointCloudViewer(QOpenGLWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.point_cloud = None  # (M,3) 点云坐标
        self.point_colors = None  # (M,3) 点云颜色（RGB，0~1）
        self.mesh = None          # Open3D 网格对象（重建结果）
        self.rotation_x = 0
        self.rotation_y = 0
        self.zoom = 5.0
        self.last_mouse_pos = None

        self.setMinimumSize(300, 300)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

    def initializeGL(self):
        glClearColor(0.1, 0.1, 0.1, 1.0)
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_PROGRAM_POINT_SIZE)
        glEnable(GL_POINT_SMOOTH)
        glEnable(GL_LIGHTING)  # 启用光照（提升网格可视化效果）
        glEnable(GL_LIGHT0)
        glEnable(GL_COLOR_MATERIAL)

        # 设置光照参数
        light_pos = [1.0, 1.0, 1.0, 0.0]  # 方向光
        glLightfv(GL_LIGHT0, GL_POSITION, light_pos)
        glLightfv(GL_LIGHT0, GL_AMBIENT, [0.3, 0.3, 0.3, 1.0])
        glLightfv(GL_LIGHT0, GL_DIFFUSE, [0.7, 0.7, 0.7, 1.0])

    def resizeGL(self, width, height):
        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45, (width / height), 0.1, 100.0)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

    def initializeGL(self):
        # 保持你的黑色环境底色（不修改）
        glClearColor(0.0, 0.0, 0.0, 1.0)  # 纯黑色背景
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_PROGRAM_POINT_SIZE)
        glEnable(GL_POINT_SMOOTH)
        glEnable(GL_LIGHTING)  # 启用光照（关键：增强对比度）
        glEnable(GL_LIGHT0)
        glEnable(GL_COLOR_MATERIAL)
        glEnable(GL_NORMALIZE)  # 归一化法线，确保光照均匀（新增）

        # 优化光照参数（黑色背景专用：增强亮度和对比度）
        light_pos = [3.0, 3.0, 3.0, 0.0]  # 更强的斜向光
        glLightfv(GL_LIGHT0, GL_POSITION, light_pos)
        glLightfv(GL_LIGHT0, GL_AMBIENT, [0.5, 0.5, 0.5, 1.0])  # 提高环境光（避免过暗）
        glLightfv(GL_LIGHT0, GL_DIFFUSE, [1.0, 1.0, 1.0, 1.0])  # 全强度漫反射（更亮）
        glLightfv(GL_LIGHT0, GL_SPECULAR, [0.8, 0.8, 0.8, 1.0])  # 新增镜面光（增强立体感）

    def paintGL(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()

        # 相机位置（保持不变）
        gluLookAt(0, 0, self.zoom,
                0, 0, 0,
                0, 1, 0)
        glRotatef(self.rotation_x, 1, 0, 0)
        glRotatef(self.rotation_y, 0, 1, 0)

        # 绘制点云（保持不变）
        if self.point_cloud is not None:
            glDisable(GL_LIGHTING)
            glPointSize(2.0)
            glBegin(GL_POINTS)
            for i, point in enumerate(self.point_cloud):
                if self.point_colors is not None and i < len(self.point_colors):
                    color = self.point_colors[i]
                    glColor3f(color[0], color[1], color[2])
                else:
                    glColor3f(0.4, 0.7, 1.0)  # 点云默认天蓝色（黑色背景下清晰）
                glVertex3f(point[0], point[1], point[2])
            glEnd()
            glEnable(GL_LIGHTING)

        # ==============================================
        # 黑色背景专用：优化网格颜色和光照
        # ==============================================
        if self.mesh is not None:
            vertices = np.asarray(self.mesh.vertices)
            triangles = np.asarray(self.mesh.triangles)
            normals = np.asarray(self.mesh.vertex_normals)

            # 中心化和缩放（保持不变）
            centroid = np.mean(vertices, axis=0)
            vertices = vertices - centroid
            max_range = np.max(np.linalg.norm(vertices, axis=1))
            if max_range > 0:
                vertices = vertices / max_range * 2.0

            # 关键：提高网格材质饱和度（黑色背景下更鲜艳）
            # 材质底色+漫反射：用明亮的青蓝色（替代之前的淡蓝，不暗沉）
            glMaterialfv(GL_FRONT_AND_BACK, GL_AMBIENT, [0.3, 0.7, 1.0, 1.0])  # 底色更亮
            glMaterialfv(GL_FRONT_AND_BACK, GL_DIFFUSE, [0.6, 1.0, 1.0, 1.0])  # 漫反射色（青蓝，饱和度高）
            glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR, [1.0, 1.0, 1.0, 1.0])  # 镜面光（高亮反光）
            glMaterialf(GL_FRONT_AND_BACK, GL_SHININESS, 30.0)  # 镜面反光强度（0-128）

            # 绘制带法线的网格（保持逻辑不变，颜色已优化）
            glBegin(GL_TRIANGLES)
            for triangle in triangles:
                for idx in triangle:
                    vertex = vertices[idx]
                    normal = normals[idx] if len(normals) > idx else [0, 0, 1]
                    normal = normal / np.linalg.norm(normal) if np.linalg.norm(normal) > 0 else [0, 0, 1]
                    glNormal3f(normal[0], normal[1], normal[2])
                    glVertex3f(vertex[0], vertex[1], vertex[2])
            glEnd()

            # 线框优化（黑色背景下更清晰，不突兀）
            glColor3f(0.2, 0.4, 0.6)  # 线框色：深青蓝（替代深灰，和网格呼应）
            glLineWidth(0.8)
            glDisable(GL_LIGHTING)
            glBegin(GL_LINES)
            for triangle in triangles:
                for i in range(3):
                    v1 = vertices[triangle[i]]
                    v2 = vertices[triangle[(i+1)%3]]
                    glVertex3f(v1[0], v1[1], v1[2])
                    glVertex3f(v2[0], v2[1], v2[2])
            glEnd()
            glEnable(GL_LIGHTING)

    def set_point_cloud(self, point_cloud_data, point_colors=None):
        """设置点云数据（支持彩色）"""
        self.clear()  # 清除原有网格
        if point_cloud_data is not None:
            # 中心化和缩放
            centroid = np.mean(point_cloud_data, axis=0)
            point_cloud_data = point_cloud_data - centroid
            max_range = np.max(np.linalg.norm(point_cloud_data, axis=1))
            if max_range > 0:
                point_cloud_data = point_cloud_data / max_range * 2.0

        self.point_cloud = point_cloud_data
        self.point_colors = point_colors
        self.rotation_x = 0
        self.rotation_y = 0
        self.zoom = 5.0
        self.update()

    def set_mesh(self, mesh):
        """设置网格数据（重建结果）"""
        self.clear()  # 清除原有点云
        self.mesh = mesh
        self.rotation_x = 0
        self.rotation_y = 0
        self.zoom = 5.0
        self.update()

    def clear(self):
        """清空点云、颜色和网格"""
        self.point_cloud = None
        self.point_colors = None
        self.mesh = None
        self.update()

    def mousePressEvent(self, event):
        """鼠标按下事件：记录位置"""
        self.last_mouse_pos = event.pos()

    def mouseMoveEvent(self, event):
        """鼠标移动事件：旋转点云"""
        if event.buttons() == Qt.LeftButton and self.last_mouse_pos:
            dx = event.x() - self.last_mouse_pos.x()
            dy = event.y() - self.last_mouse_pos.y()
            self.rotation_y += dx * 0.5
            self.rotation_x += dy * 0.5
            self.last_mouse_pos = event.pos()
            self.update()

    def wheelEvent(self, event):
        """鼠标滚轮事件：缩放点云"""
        delta = event.angleDelta().y()
        if delta > 0:
            self.zoom -= 0.5  # 放大
        else:
            self.zoom += 0.5  # 缩小
        # 限制缩放范围（1.0~20.0）
        self.zoom = max(1.0, min(self.zoom, 20.0))
        self.update()